// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/array/array_run_end.h>
#include <nebula/array/builder_primitive.h>
#include <nebula/array/util.h>
#include <turbo/log/logging.h>
#include <nebula/util/ree_util.h>

namespace nebula {

    // ----------------------------------------------------------------------
    // RunEndEncodedArray

    RunEndEncodedArray::RunEndEncodedArray(const std::shared_ptr<ArrayData> &data) {
        this->set_data(data);
    }

    RunEndEncodedArray::RunEndEncodedArray(const std::shared_ptr<DataType> &type,
                                           int64_t length,
                                           const std::shared_ptr<Array> &run_ends,
                                           const std::shared_ptr<Array> &values,
                                           int64_t offset) {
        this->set_data(ArrayData::create(type, length,
                /*buffers=*/{nullptr},
                /*child_data=*/{run_ends->data(), values->data()},
                /*null_count=*/0, offset));
    }

    turbo::Result<std::shared_ptr<RunEndEncodedArray>> RunEndEncodedArray::create(
            const std::shared_ptr<DataType> &type, int64_t logical_length,
            const std::shared_ptr<Array> &run_ends, const std::shared_ptr<Array> &values,
            int64_t logical_offset) {
        if (type->id() != Type::RUN_END_ENCODED) {
            return turbo::invalid_argument_error("Type must be RUN_END_ENCODED");
        }
        const auto *ree_type = turbo::checked_cast<const RunEndEncodedType *>(type.get());
        TURBO_RETURN_NOT_OK(ree_util::ValidateRunEndEncodedChildren(
                *ree_type, logical_length, run_ends->data(), values->data(), 0, logical_offset));
        return std::make_shared<RunEndEncodedArray>(type, logical_length, run_ends, values,
                                                    logical_offset);
    }

    turbo::Result<std::shared_ptr<RunEndEncodedArray>> RunEndEncodedArray::create(
            int64_t logical_length, const std::shared_ptr<Array> &run_ends,
            const std::shared_ptr<Array> &values, int64_t logical_offset) {
        auto run_end_type = run_ends->type();
        auto value_type = values->type();
        if (!RunEndEncodedType::RunEndTypeValid(*run_end_type)) {
            return turbo::invalid_argument_error("Run end type must be int16, int32 or int64");
        }
        auto ree_type = run_end_encoded(std::move(run_end_type), std::move(value_type));
        return create(ree_type, logical_length, run_ends, values, logical_offset);
    }

    void RunEndEncodedArray::set_data(const std::shared_ptr<ArrayData> &data) {
        KCHECK_EQ(data->type->id(), Type::RUN_END_ENCODED);
        const auto *ree_type =
                turbo::checked_cast<const RunEndEncodedType *>(data->type.get());
        KCHECK_EQ(data->child_data.size(), 2);
        KCHECK_EQ(ree_type->run_end_type()->id(), data->child_data[0]->type->id());
        KCHECK_EQ(ree_type->get_value_type()->id(), data->child_data[1]->type->id());

        Array::set_data(data);
        run_ends_array_ = make_array(this->data()->child_data[0]);
        values_array_ = make_array(this->data()->child_data[1]);
    }

    namespace {

        template<typename RunEndType>
        turbo::Result<std::shared_ptr<Array>> MakeLogicalRunEnds(const RunEndEncodedArray &self,
                                                                 MemoryPool *pool) {
            using RunEndCType = typename RunEndType::c_type;
            if (self.offset() == 0) {
                const auto &run_ends = *self.run_ends();
                if (self.length() == 0) {
                    return run_ends.slice(0, 0);
                }

                // If offset==0 and the non-zero logical length aligns perfectly with a
                // physical run-end, we can return a slice of the run-ends array.
                const int64_t physical_length = self.FindPhysicalLength();
                const auto *run_end_values = self.data()->child_data[0]->get_values<RunEndCType>(1);
                if (run_end_values[physical_length - 1] == self.length()) {
                    return run_ends.slice(0, physical_length);
                }

                // Otherwise we need to copy the run-ends array and adjust only the very
                // last run-end.
                auto new_run_ends_data = ArrayData::create(run_ends.type(), physical_length, 0, 0);
                {
                    TURBO_MOVE_OR_RAISE(auto buffer,
                                        allocate_buffer(physical_length * sizeof(RunEndCType), pool));
                    new_run_ends_data->buffers = {nullptr, std::move(buffer)};
                }
                auto *new_run_end_values = new_run_ends_data->GetMutableValues<RunEndCType>(1);
                memcpy(new_run_end_values, run_end_values,
                       (physical_length - 1) * sizeof(RunEndCType));
                new_run_end_values[physical_length - 1] = static_cast<RunEndCType>(self.length());
                return make_array(std::move(new_run_ends_data));
            }

            // When the logical offset is non-zero, all run-end values need to be adjusted.
            int64_t physical_offset = self.FindPhysicalOffset();
            int64_t physical_length = self.FindPhysicalLength();

            const auto *run_end_values = self.data()->child_data[0]->get_values<RunEndCType>(1);
            NumericBuilder<RunEndType> builder(pool);
            TURBO_RETURN_NOT_OK(builder.resize(physical_length));
            if (physical_length > 0) {
                for (int64_t i = 0; i < physical_length - 1; i++) {
                    const auto run_end = run_end_values[physical_offset + i] - self.offset();
                            DKCHECK_LT(run_end, self.length());
                    TURBO_RETURN_NOT_OK(builder.append(static_cast<RunEndCType>(run_end)));
                }
                        DKCHECK_GE(run_end_values[physical_offset + physical_length - 1] - self.offset(),
                                         self.length());
                TURBO_RETURN_NOT_OK(builder.append(static_cast<RunEndCType>(self.length())));
            }
            return builder.finish();
        }

    }  // namespace

    turbo::Result<std::shared_ptr<Array>> RunEndEncodedArray::LogicalRunEnds(
            MemoryPool *pool) const {
                DKCHECK(data()->child_data[0]->buffers[1]->is_cpu());
        switch (run_ends_array_->type_id()) {
            case Type::INT16:
                return MakeLogicalRunEnds<Int16Type>(*this, pool);
            case Type::INT32:
                return MakeLogicalRunEnds<Int32Type>(*this, pool);
            default:
                        DKCHECK_EQ(run_ends_array_->type_id(), Type::INT64);
                return MakeLogicalRunEnds<Int64Type>(*this, pool);
        }
    }

    std::shared_ptr<Array> RunEndEncodedArray::LogicalValues() const {
        const int64_t physical_offset = FindPhysicalOffset();
        const int64_t physical_length = FindPhysicalLength();
        return make_array(data()->child_data[1]->slice(physical_offset, physical_length));
    }

    int64_t RunEndEncodedArray::FindPhysicalOffset() const {
        const ArraySpan span(*this->data_);
        return ree_util::FindPhysicalIndex(span, 0, span.offset);
    }

    int64_t RunEndEncodedArray::FindPhysicalLength() const {
        const ArraySpan span(*this->data_);
        return ree_util::FindPhysicalLength(span);
    }

}  // namespace nebula
